package main

import (
	"fmt"
	"io"
	"net"
	"strconv"
)

type ssuint8 uint8

const (
	// 版本
	VERSION5 = ssuint8(0x05)
	// method
	METHOD_NO_AUTH  = ssuint8(0x00)
	METHOD_USER_PWD = ssuint8(0x02)

	// cmd
	CMD_CONNECT = ssuint8(0x01)
	CMD_BIND = ssuint8(0x02)
	CMD_UDP = ssuint8(0x03)

	// 地址类型
	ADDR_IP4 = ssuint8(0x01)
	ADDR_IP6 = ssuint8(0x04)
	ADDR_DOMAIN = ssuint8(0x03)

	// request reply

)
// 认证方式请求
type MethodsRequest struct {
	version ssuint8
	methodsCount ssuint8
	methods []byte

	conn *net.Conn
}

// 协议方式请求
type ProtocolRequest struct {
	// 版本
	version ssuint8
	// 要做什么：
	cmd ssuint8
	// 保留位
	rsv ssuint8
	// 地址类型
	addrType ssuint8

	// 目的地址
	dstHost string
	// 目的地址端口
	dstPort string

	// request接收的数据缓存
	buf []byte

	conn *net.Conn
}

// 服务
type Service struct {
	conn *net.Conn

	dstHost string
	dstPort string

	// tcp, udp
	network string

	buf []byte
}

func main() {
	Socks5(":12390")
}

func Socks5(addr string)  {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer func() {
		_ = listener.Close()
	}()

	// 接受连接
	for {
		conn, err := listener.Accept()
		if err != nil {
			_ = conn.Close()
			panic(err)
		}

		// 接收、发送数据
		go func() {
			defer func() {
				//err = recover().(error)
			}()

			var service Service
			// 握手
			if err := service.handshake(conn); err != nil  {
				panic(err)
			}


			// 开始转发
			switch service.network {
			case "tcp":
				err = service.proxyTCP()
			case "udp":
				err = service.proxyUDP()
			default:
				panic("unknown network" + service.network)
			}
			if err != nil {
				panic(err)
			}
		}()
	}
}

func(s *Service) handshake(conn net.Conn)(err error) {

	var mReq MethodsRequest
	var req ProtocolRequest

	// 第一步，授权鉴定
	if mReq, err = s.newMethodRequest(&conn); err == nil {
		if err = mReq.reply(); err != nil {
			return
		}
	}

	// 第二步，确认使用什么协议
	if req, err = s.newRequest(&conn); err == nil {
		if err = req.reply(); err != nil {
			return
		}
		s.buf = req.buf
		fmt.Println(string(s.buf))
	}
	return
}

func(s *Service) newMethodRequest(conn *net.Conn) (mReq MethodsRequest, err error) {
	mReq.version = VERSION5
	mReq.conn = conn

	buf := make([]byte, 512)
	n, err := (*conn).Read(buf)
	if n > 0 {
		fmt.Println(buf[:n])
		if ssuint8(buf[0]) != VERSION5 {
			err = fmt.Errorf("not support version: %x", buf[0])
			return
		}
		mReq.methodsCount = ssuint8(buf[1])
		mReq.methods = buf[2:n]
	}
	return mReq, err
}


func(mReq *MethodsRequest) reply()(err error) {
	// reply

	// no user
	_, err = (*mReq.conn).Write([]byte{byte(mReq.version), byte(METHOD_NO_AUTH)})

	//// must user/password
	//_, err = (*mReq.conn).Write([]byte{byte(mReq.version), byte(METHOD_USER_PWD)})
	//
	//buf := make([]byte, 1024)
	//n, err := (*mReq.conn).Read(buf)
	//if n > 0 {
	//	fmt.Println(string(buf[:n]))
	//}
	//
	//_, err = (*mReq.conn).Write([]byte{byte(mReq.version), byte(METHOD_NO_AUTH)})

	return
}

// new request
func(s *Service) newRequest(conn *net.Conn)(req ProtocolRequest, err error)  {
	req.conn = conn
	req.version = VERSION5

	buf := make([]byte, 512)
	n, err := (*conn).Read(buf)
	if n > 0 {
		fmt.Println(buf[:n])

		req.buf = buf[:n]
		req.version = ssuint8(buf[0])
		req.cmd = ssuint8(buf[1])
		req.rsv = ssuint8(buf[2])
		req.addrType = ssuint8(buf[3])

		switch req.addrType {
		case ADDR_IP4:
			req.dstHost = net.IPv4(buf[4], buf[5], buf[6], buf[7]).String()
		case ADDR_IP6:
			req.dstHost = net.IP{buf[4], buf[5], buf[6], buf[7], buf[8], buf[9], buf[10], buf[11], buf[12], buf[13], buf[14], buf[15], buf[16], buf[17], buf[18], buf[19]}.String()
		case ADDR_DOMAIN:
			// 第4位为域名长度，第5位到倒数第2位为域名
			req.dstHost = string(buf[5:n-2])
		}

		// 端口为最后两个字节
		req.dstPort = strconv.Itoa(int(buf[n-2]) << 8 | int(buf[n-1]))
	}

	if err != nil {
		return
	}

	s.conn = conn
	s.dstHost = req.dstHost
	s.dstPort = req.dstPort
	switch req.cmd {
	case CMD_CONNECT:
		s.network = "tcp"
	case CMD_UDP:
		s.network = "udp"
	}

	return req, err
}

// request reply
func(req *ProtocolRequest) reply()(err error) {
	switch req.cmd {
	case CMD_CONNECT:
		req.buf[1] = byte(ssuint8(0x00))
		fmt.Println(req.buf[1])
	case CMD_UDP:
		req.buf[1] = byte(ssuint8(0x00))
	case CMD_BIND:
		req.buf[1] = byte(ssuint8(0x09))
	}
	_, err = (*req.conn).Write(req.buf)
	return
}

func(s *Service) proxyTCP() error {
	addr := net.JoinHostPort(s.dstHost, s.dstPort)
	fmt.Println(addr)
	dstConn, err := net.Dial(s.network, addr)
	if err != nil {
		return err
	}

	defer func() {
		_ = dstConn.Close()
	}()


	if err == nil {
		errCh := make(chan error, 2)
		go proxy(*s.conn, dstConn, errCh)
		go proxy(dstConn, *s.conn, errCh)

		for i := 0; i < 2; i++ {
			err = <-errCh
			if err != nil {
				return err
			}
		}
	}
	return err
}

func(s *Service) proxyUDP() error {
	return nil
}


func proxy(dst io.Writer, src io.Reader, errCh chan error) {
	_, err := io.Copy(dst, src)
	errCh <- err
}